// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{collections::{hash_map::Entry, HashMap}, process::{ExitStatus, Stdio}, str::FromStr, sync::Arc};

use futures::{stream::SplitSink, SinkExt};
use lazy_static::lazy_static;
use log::{error, trace, warn};
use maestro_derive::MsgTag;
use maestro_message::serialize_msg;
use ohmw_core::bincode::Encode;
use ohmw_derive::{BorrowDecode, Encode};
use tokio::{io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}, net::TcpStream, process::{Child, Command}, sync::Mutex, task::JoinHandle};
use tokio_tungstenite::{tungstenite::Message, WebSocketStream};
use tokio_util::sync::CancellationToken;

lazy_static! {
    static ref NODE_MAP: Arc<Mutex<HashMap<String, (JoinHandle<()>, CancellationToken)>>> = Arc::new(Mutex::new(HashMap::new()));
}

#[derive(BorrowDecode)]
pub struct CustomNodeOperateReq<'a> {
    pub node_name: &'a str,
    pub token: &'a str,
    pub operate: &'a str,
    pub path: &'a str,
    pub code: &'a str,
}

pub enum CustomNodeOperate {
    Start,
    Stop,
    Reload,
}

#[repr(u8)]
#[derive(Clone, Copy)]
pub enum CustomNodeOperateStatus {
    Running,
    Stopped,
    Out,
    Err,
}

#[derive(Encode, MsgTag)]
#[msg_tag(case = "snake_case")]
struct CustomNodeOperateRet<'a> {
    token: &'a str,
    status: CustomNodeOperateStatus,
    log: String,
}

impl Encode for CustomNodeOperateStatus {
    unsafe fn encode_unchecked<E: ohmw_core::bincode::Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
        unsafe { self.as_str().encode_unchecked(encoder, buf) }
    }
    
    fn encode<E: ohmw_core::bincode::Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
        self.as_str().encode(encoder, buf)
    }
}

impl FromStr for CustomNodeOperate {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "start" => Ok(CustomNodeOperate::Start),
            "stop" => Ok(CustomNodeOperate::Stop),
            "reload" => Ok(CustomNodeOperate::Reload),
            _ => Err(anyhow::anyhow!("Invalid CustomNodeOperate: {}", s)),
        }
    }
}

impl CustomNodeOperateStatus {
    fn as_str(&self) -> &'static str {
        match self {
            CustomNodeOperateStatus::Running => "running",
            CustomNodeOperateStatus::Stopped => "stopped",
            CustomNodeOperateStatus::Out => "out",
            CustomNodeOperateStatus::Err => "err",
        }
    }
}

async fn forward_output<R: AsyncReadExt + std::marker::Unpin>(
    mut reader: BufReader<R>,
    sink: Arc<Mutex<SplitSink<WebSocketStream<TcpStream>, Message>>>,
    token: String,
    status: CustomNodeOperateStatus,
) {
    let mut buf = String::new();

    loop {
        match reader.read_line(&mut buf).await {
            Ok(n) => {
                if n == 0 {
                    break;
                }

                if !buf.is_empty() {
                    let ret = CustomNodeOperateRet {
                        token: &token,
                        status,
                        log: std::mem::take(&mut buf),
                    };
                    
                    if let Err(err) = sink.lock().await.send(serialize_msg(ret).unwrap().into()).await {
                        error!("Failed to send output to client: {}", err);
                        break;
                    }
                }
            },
            Err(err) => {
                error!("Failed to read output from node: {}", err);
                break;
            },
        }
    }
}

async fn run_python_node(
    sink: Arc<Mutex<SplitSink<WebSocketStream<TcpStream>, Message>>>,
    node_name: &str,
    token: &str,
    path: &str,
) -> anyhow::Result<(Child, JoinHandle<()>, JoinHandle<()>)> {
    let mut child = Command::new("python3")
       .arg(path)
       .stdout(Stdio::piped())
       .stderr(Stdio::piped())
       .kill_on_drop(true)
       .spawn()?;
    let stdout = child.stdout.take().unwrap();
    let stderr = child.stderr.take().unwrap();
    let stdout_reader = tokio::io::BufReader::new(stdout);
    let stderr_reader = tokio::io::BufReader::new(stderr);

    let token_clone = token.to_string();
    let token_clone2 = token.to_string();
    let thread_stdout = tokio::spawn(forward_output(
        stdout_reader,
        sink.clone(),
        token_clone,
        CustomNodeOperateStatus::Out,
    ));
    let thread_stderr = tokio::spawn(forward_output(
        stderr_reader,
        sink.clone(),
        token_clone2,
        CustomNodeOperateStatus::Err,
    ));

    trace!("Node {} with token {} started", node_name, token);
    let ret = CustomNodeOperateRet {
        token,
        status: CustomNodeOperateStatus::Running,
        log: "running".to_string(),
    };
    if let Err(err) = sink.lock().await.send(serialize_msg(ret).unwrap().into()).await {
        error!("Failed to send output to client: {}", err);
    }
    
    Ok((child, thread_stdout, thread_stderr))
}

async fn wait_python_node(
    status: Result<ExitStatus, tokio::io::Error>,
    thread_stdout: JoinHandle<()>,
    thread_stderr: JoinHandle<()>,
    sink: Arc<Mutex<SplitSink<WebSocketStream<TcpStream>, Message>>>,
    node_name: &str,
    token: &str,
) -> anyhow::Result<()> {
    let status = status?;
    if !status.success() {
        error!("Node {} with token {} exited with status {}", node_name, token, status.code().unwrap());
    }

    thread_stdout.await?;
    thread_stderr.await?;

    trace!("Node {} with token {} stopped", node_name, token);
    let ret = CustomNodeOperateRet {
        token: &token,
        status: CustomNodeOperateStatus::Stopped,
        log: "stopped".to_string(),
    };
    if let Err(err) = sink.lock().await.send(serialize_msg(ret).unwrap().into()).await {
        error!("Failed to send output to client: {}", err);
    }

    Ok(())
}

async fn start_node(
    sink: Arc<Mutex<SplitSink<WebSocketStream<TcpStream>, Message>>>,
    node_name: &str,
    token: &str,
    path: &str,
    code: &str,
) -> Option<(JoinHandle<()>, CancellationToken)> {
    if let Err(err) = tokio::fs::write(path, code).await {
        error!("Failed to write code to file at {}: {}", path, err);
        return None;
    }

    let cancel_token = CancellationToken::new();
    let cancel_token_clone = cancel_token.clone();

    let sink = sink.clone();
    let node_name = node_name.to_string();
    let token = token.to_string();
    let path = path.to_string();
    let thread = tokio::spawn(async move {
        let (mut child, thread_stdout, thread_stderr) = match run_python_node(
            sink.clone(),
            &node_name,
            &token,
            &path,
        ).await {
            Ok(res) => res,
            Err(err) => {
                error!("Failed to run node {} with token {}: {}", node_name, token, err);
                return;
            },
        };

        tokio::select! {
            _ = cancel_token_clone.cancelled() => {
                if let Err(err) = child.kill().await {
                    error!("Failed to kill node {} with token {}: {}", node_name, token, err);
                } else {
                    trace!("Node {} with token {} killed", node_name, token);
                }
            },
            status = child.wait() => {
                if let Err(err) = wait_python_node(status, thread_stdout, thread_stderr, sink, &node_name, &token).await {
                    error!("Failed to wait for node {} with token {}: {}", node_name, token, err);
                }
            },
        }
    });

    Some((thread, cancel_token))
}

async fn stop_node(
    thread: &mut JoinHandle<()>,
    cancel_token: CancellationToken,
    sink: Arc<Mutex<SplitSink<WebSocketStream<TcpStream>, Message>>>,
    node_name: &str,
    token: &str,
) {
    if thread.is_finished() {
        return;
    }

    cancel_token.cancel();
    let _ = thread.await;

    if let Err(err) = cleanup_dead_nodes().await {
        error!("Failed to cleanup dead nodes: {}", err);
    }

    trace!("Node {} with token {} stopped", node_name, token);
    let ret = CustomNodeOperateRet {
        token: token,
        status: CustomNodeOperateStatus::Stopped,
        log: "stopped".to_string(),
    };
    if let Err(err) = sink.lock().await.send(serialize_msg(ret).unwrap().into()).await {
        error!("Failed to send output to client: {}", err);
    }
}

async fn cleanup_dead_nodes() -> anyhow::Result<()> {
    // TODO: use rust api to register the custom node only?
    let mut rosnode = Command::new("rosnode")
        .arg("cleanup")
        .stdin(Stdio::piped())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()?;
    rosnode.stdin.take().unwrap().write_all(b"y").await?;
    rosnode.wait().await?;
    Ok(())
}

pub async fn custom_node_operate(
    sink: Arc<Mutex<SplitSink<WebSocketStream<TcpStream>, Message>>>,
    node_name: &str,
    operate: CustomNodeOperate,
    token: &str,
    path: &str,
    code: &str,
) {
    match operate {
        CustomNodeOperate::Start => {
            match NODE_MAP.lock().await.entry(token.to_string()) {
                Entry::Occupied(mut occupied_entry) => {
                    if occupied_entry.get().0.is_finished() {
                        if let Some(res) = start_node(sink, node_name, token, path, code).await {
                            *occupied_entry.get_mut() = res;
                        }
                    } else {
                        warn!("Node {} with token {} is already running, reloading", node_name, token);
                        let cancel_token = occupied_entry.get().1.clone();
                        stop_node(&mut occupied_entry.get_mut().0, cancel_token, sink.clone(), node_name, token).await;
                        if let Some(thread) = start_node(sink, node_name, token, path, code).await {
                            *occupied_entry.get_mut() = thread;
                        }
                    }
                },
                Entry::Vacant(vacant_entry) => {
                    if let Some(thread) = start_node(sink, node_name, token, path, code).await {
                        vacant_entry.insert(thread);
                    }
                },
            }
        },
        CustomNodeOperate::Stop => {
            if let Some((mut thread, cancel_token)) = NODE_MAP.lock().await.remove(token) {
                stop_node(&mut thread, cancel_token, sink, node_name, token).await;
            } else {
                warn!("Node {} with token {} is not running", node_name, token);
            }
        },
        CustomNodeOperate::Reload => {
            match NODE_MAP.lock().await.entry(token.to_string()) {
                Entry::Occupied(mut occupied_entry) => {
                    let cancel_token = occupied_entry.get().1.clone();
                    stop_node(&mut occupied_entry.get_mut().0, cancel_token, sink.clone(), node_name, token).await;
                    if let Some(thread) = start_node(sink, node_name, token, path, code).await {
                        *occupied_entry.get_mut() = thread;
                    }
                },
                Entry::Vacant(vacant_entry) => {
                    warn!("Node {} with token {} is not running", node_name, token);
                    if let Some(thread) = start_node(sink, node_name, token, path, code).await {
                        vacant_entry.insert(thread);
                    }
                },
            }
        },
    }
}

pub async fn stop_custom_nodes() {
    let mut node_map = NODE_MAP.lock().await;
    for (_, (thread, cancel_token)) in node_map.drain() {
        cancel_token.cancel();
        let _ = thread.await;

        if let Err(err) = cleanup_dead_nodes().await {
            error!("Failed to cleanup dead nodes: {}", err);
        }
    }
}
